# coding=utf-8
import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

# Data Visualization
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 unused import
import matplotlib.pyplot as plt
import seaborn as sns

# Other packages
import warnings
warnings.filterwarnings("ignore")

SP_df = pd.read_csv("exams_students_100.csv")

def data_exploration():
    # 数学分数数据特征
    fig, ax = plt.subplots(ncols=2,figsize=(15,6))
    sns.distplot(SP_df['math score'],
                ax=ax[1])
    sns.boxplot(y=SP_df['math score'],
            ax=ax[0])
    # 阅读分数数据特征
    fig, ax = plt.subplots(ncols=2,figsize=(15,6))
    sns.distplot(SP_df['reading score'],
                ax=ax[1])
    sns.boxplot(y=SP_df['reading score'],
            ax=ax[0])
    # 写作分数数据特征
    fig, ax = plt.subplots(ncols=2,figsize=(15,6))
    sns.distplot(SP_df['writing score'],
                ax=ax[1])
    sns.boxplot(y=SP_df['writing score'],
            ax=ax[0])

    # 影响属性对各科分数的特征分析
    # 从 test preparation course维度下
    g = sns.pairplot(data=SP_df,
                hue='test preparation course'
                )
    g.fig.set_size_inches(15,8)
    # 从 gender维度下
    g = sns.pairplot(data=SP_df,
                hue='gender'
                )
    g.fig.set_size_inches(15,8)
    # 从 lunch 维度下
    g = sns.pairplot(data=SP_df,
                hue='lunch'
                )
    g.fig.set_size_inches(15,8)

    def autolabel(rects, xpos='center'):
        """
        Attach a text label above each bar in *rects*, displaying its height.

        *xpos* indicates which side to place the text w.r.t. the center of
        the bar. It can be one of the following {'center', 'right', 'left'}.
        """

        xpos = xpos.lower()  # normalize the case of the parameter
        ha = {'center': 'center', 'right': 'left', 'left': 'right'}
        offset = {'center': 0.5, 'right': 0.57, 'left': 0.43}  # x_txt = x + w*off

        for rect in rects:
            height = rect.get_height()
            ax.text(rect.get_x() + rect.get_width()*offset[xpos], 1.01*height,
                    '{}'.format(height), ha=ha[xpos], va='bottom')
    
    data = SP_df.copy()
    # 从种族角度看各科成绩
    # 对总的平均分的关联表现
    race = round(data.groupby(by = data['race/ethnicity']).mean(), 1)
    x = list(race.index)
    y = round(race.mean(axis=1),1)

    fig, ax = plt.subplots(figsize=(12,6))

    rects = ax.bar(x, y)

    ax.set_ylabel('Average Scores')
    ax.set_title('Average scores by race/ethnicity')
    ax.set_xticklabels((x))

    autolabel(rects, 'center')
    # 种族对各科成绩的表现
    x = race.index
    y1 = race['math score']
    y2 = race['reading score']
    y3 = race['writing score']

    ind = np.arange(len(x))  # the x locations for the groups
    width = 0.35  # the width of the bars

    fig, ax = plt.subplots(figsize=(12, 6))
    rects1 = ax.bar(ind - width/2, y1, width/2,
                color='SkyBlue', label='Math')
    rects2 = ax.bar(ind, y2, width/2,
                color='IndianRed', label='Reading')
    rects3 = ax.bar(ind + width/2, y3, width/2,
                color='Green', label='Writing')

    # Add some text for labels, title and custom x-axis tick labels, etc.
    ax.set_ylabel('Average Scores')
    ax.set_title('Average scores by race/ethnicity')
    ax.set_xticks(ind)
    ax.set_xticklabels((x))
    ax.legend(loc=3)

    autolabel(rects1, "center")
    autolabel(rects2, "center")
    autolabel(rects3, "center")

    # 父母受教育水平对其学科成绩的影响
    # 对学生总的平均分的表现
    parent_edu = round(SP_df.groupby(by = SP_df['parental level of education']).mean(),1)
    x = list(parent_edu.index)
    y = round(parent_edu.mean(axis=1),1)

    fig, ax = plt.subplots(figsize=(12,6))

    rects = ax.bar(x, y)

    ax.set_ylabel('Average Scores')
    ax.set_title('Average scores by parent\'s education level')
    ax.set_xticklabels((x))
    autolabel(rects, 'center')
    # 父母受教育水平对各科成绩的表现
    x = parent_edu.index
    y1 = parent_edu['math score']
    y2 = parent_edu['reading score']
    y3 = parent_edu['writing score']

    ind = np.arange(6)  # the x locations for the groups
    width = 0.35  # the width of the bars

    fig, ax = plt.subplots(figsize=(12, 6))
    rects1 = ax.bar(ind - width/2, y1, width/2,
                color='SkyBlue', label='Math')
    rects2 = ax.bar(ind, y2, width/2,
                color='IndianRed', label='Reading')
    rects3 = ax.bar(ind + width/2, y3, width/2,
                color='Green', label='Writing')

    # Add some text for labels, title and custom x-axis tick labels, etc.
    ax.set_ylabel('Average Scores')
    ax.set_title('Average scores by subject and parent\'s education level')
    ax.set_xticks(ind)
    ax.set_xticklabels((x))
    ax.legend(loc=3)

    autolabel(rects1, "center")
    autolabel(rects2, "center")
    autolabel(rects3, "center")

    # lunch与preparation course
    x = ['Test prep', 'No test prep']
    y = [data[(data['lunch'] == 'free/reduced') & (data['test preparation course'] == 'completed')]['lunch'].count(),
    data[(data['lunch'] == 'free/reduced') & (data['test preparation course'] == 'none')]['lunch'].count()]

    fig, ax = plt.subplots(figsize=(12,6))

    rects = ax.bar(x, y)

    ax.set_ylabel('Count')
    ax.set_title('Students studying in schools that provide free/reduced price lunch by test prep course')
    ax.set_xticklabels((x))

    autolabel(rects, 'center')

    plt.show()


# 画出数据的3-D图
def plot_students_performance():
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    xs = SP_df['math score']
    ys = SP_df['reading score']
    zs = SP_df['writing score']
    ax.scatter(xs, ys, zs)

    ax.set_xlabel('math score Label')
    ax.set_ylabel('reading score Label')
    ax.set_zlabel('writing score Label')

    plt.show()

def scores_regression():
    N = SP_df.shape[0]
    X_matrix = SP_df.head(N)[['math score', 'reading score']].values
    X_matrix = np.matrix(np.c_[np.ones(N), X_matrix])
    Y_matrix = SP_df.head(N)[['writing score']].values
    Y_matrix = np.matrix(Y_matrix)

    print('X_matrix shape: ', X_matrix.shape)
    print('Y_matrix shape: ', Y_matrix.shape)

    beta_hat =  np.linalg.inv((X_matrix.T) * X_matrix) * (X_matrix.T) * Y_matrix

    X_inverse = (X_matrix.T * X_matrix).I

    return beta_hat, X_inverse, N

# 预测检验
def predict_test(beta_hat, x_inverse, N_train=0):
    SP_test = pd.read_csv("exams_students_10.csv")
    N = SP_test.shape[0]
    X_matrix = SP_test.head(N)[['math score', 'reading score']].values
    X_matrix_test = np.matrix(np.c_[np.ones(N), X_matrix])
    Y_matrix = SP_test.head(N)[['writing score']].values
    Y_matrix_test = np.matrix(Y_matrix)

    Y_test_hat = X_matrix_test * beta_hat
    # 计算Q值
    Q = (Y_matrix_test - Y_test_hat).T * (Y_matrix_test - Y_test_hat)
    # 计算sigma hat的值
    sigma_hat = np.sqrt(Q/(N-2))
    # t-apha
    t_apha = 1.8595
    # 对每个测试样本进行预测，如果落在预测区间则复合预测
    y_predict_result = []
    for i in range(N):
        x_formula = np.sqrt(1+X_matrix_test[i] * x_inverse * (X_matrix_test[i]).T)
        d = t_apha * sigma_hat * x_formula
        if abs(Y_test_hat[i] - Y_matrix[i]) <= d:
            y_predict_result.append(True)
        else:
            y_predict_result.append(False)
    return y_predict_result

# 因子分析
def factor_analysis():
    pass


if __name__=="__main__":
    # 数据探索部分
    data_exploration()

    # 画出学生成绩的散点图
    plot_students_performance()

    # 计算成绩的回归方程
    beta_hat, X_inverse, N = scores_regression()
    print("beta_hat: ", beta_hat)
    predict_test_result = predict_test(beta_hat, X_inverse)
    print("predict test result: ", predict_test_result)